<template>
  <div class="train">
    <a-page-header
      style="border: 0px solid rgb(235, 237, 240)"
      title="YOLOv5 train"
      sub-title="训练用户自定义数据"
    >
    </a-page-header>
    <div class="steps">
      <a-steps :current="current" @change="onChange">
        <a-step title="Step 1" description="选择标注数据集" />
        <a-step title="Step 2" description="选择预训练权重" />
        <a-step title="Step 3" description="其他训练参数" />
        <a-step title="Step 4" description="开始训练" />
      </a-steps>
    </div>

    <div class="steps-content" v-if="current == 0">
      <p class="steps-text">
        选择已标注完成的数据集 
      </p>
      <div class="steps-input">
        <a-select
          v-model:value="dataset"
          :options="dataset_options"
          style="width: 100%; text-align: left"
          ref="select"
          @change="changeDataset()"
        >
        </a-select>
      </div>
      <p class="error" v-if="show_error">{{ error_message }}</p>
      <a-button type="primary" shape="round" class="step-next" @click="checkLables()">
        下一步
      </a-button>
    </div>
            
    <div class="steps-content" v-else-if="current == 1">
      <p class="steps-text">
        选择 YOLOv5 内置预训练权重
      </p>
      <div class="steps-input">
        <a-select
          v-model:value="weight"
          :options="weight_options"
          style="width: 100%; text-align: left"
          ref="select"
          @change="changeWeight()"
        >
        </a-select>
      </div>
      <p class="error" v-if="show_error">{{ error_message }}</p>
      <a-button type="primary" shape="round" class="step-next" @click="nextStep()">
        下一步
      </a-button>
    </div>

    <div class="steps-content" v-else-if="current == 2">
      <p class="steps-text">
        设置训练参数  
      </p>
      <div class="steps-input">
        <a-form :model="train_param" :label-col="labelCol" :wrapper-col="wrapperCol">
          <a-form-item label="参数模型">
            <a-select v-model:value="train_param.model" placeholder="yolov5s.yaml">
              <a-select-option value="yolov5s.yaml">
                yolov5s.yaml
              </a-select-option>
              <a-select-option value="yolov5m.yaml">
                yolov5m.yaml
              </a-select-option>
              <a-select-option value="yolov5l.yaml">
                yolov5l.yaml
              </a-select-option>
              <a-select-option value="yolov5x.yaml">
                yolov5x.yaml
              </a-select-option>
            </a-select>
          </a-form-item>
          <a-form-item label="Epochs">
            <a-input-number v-model:value="train_param.epochs" />
            <span style="fontSize: 13px"> * 训练迭代轮数</span>
          </a-form-item>
          <a-form-item label="Batch size">
            <a-input-number v-model:value="train_param.batch" />
            <span style="fontSize: 13px"> * 一次训练所选取的样本数</span>
          </a-form-item>
          <a-form-item label="其他参数">
            <a-input v-model:value="train_param.other" type="textarea" />
          </a-form-item>
        </a-form>
      </div>
      <p class="error" v-if="show_error">{{ error_message }}</p>
      <a-button type="primary" shape="round" class="step-next" @click="nextStep()">
        下一步
      </a-button>
    </div>

    <a-row type="flex" class="steps-content train-content" v-else-if="current == 3" justify="space-around">
      <a-col span="12">
          <a-descriptions title="训练参数">
          <a-descriptions-item label="Dataset" span="3">
            {{ dataset }}
          </a-descriptions-item>
          <a-descriptions-item label="Weight" span="3">
            {{ weight }}
          </a-descriptions-item>
          <a-descriptions-item label="Model" span="3">
            {{ train_param.model }}
          </a-descriptions-item>
          <a-descriptions-item label="Epochs" span="2">
            {{ train_param.epochs }}
          </a-descriptions-item>
          <a-descriptions-item label="Batch size" span="1">
            {{ train_param.batch }}
          </a-descriptions-item>
          <a-descriptions-item label="Other options" span="3">
            {{ train_param.other }}
          </a-descriptions-item>
          <a-descriptions-item label="最终训练代码" span="3">
          </a-descriptions-item>
        </a-descriptions>

        <a-textarea v-model:value="train_code" showCount :maxlength="100"></a-textarea>
        <a-popconfirm 
          :visible="confirm_visible"
          title="权重已存在，是否继续训练？" 
          ok-text="确定" 
          cancel-text="取消"
          @confirm="startTrain"
          @cancel="confirm_visible = !confirm_visible"
          >
          <a-button :loading="train_loading" type="primary" shape="round" class="step-next" @click="checkTrain()">
            运行
          </a-button>
        </a-popconfirm>
      </a-col>

      <a-col span="1"></a-col>
      <a-col span="11">
        <a-descriptions title="训练日志"></a-descriptions>
        <a-card :loading="log_loading" title="" id="log-content"
          style="padding: 10px; max-height: 360px; overflow: auto; white-space: pre-wrap"
        >
        {{ train_code }}
        </a-card>
      </a-col>

    </a-row>

  </div>
</template>

<script>
import { message } from 'ant-design-vue';
import qs from 'qs'
export default {
  data() {
    return {
      current: 0,  
      value: '',
      show_error: false, 
      error_message: '', 
      yolov_weights_url: this.$axios.defaults.baseURL + "/yolov/get_weights",
      weight: 'yolov5s.pt',
      weight_options: [
        // {
        //   value: 'yolov5s',
        // }
      ], 
      yolov_datasets_url: this.$axios.defaults.baseURL + "/annotate/getDatasetList",
      dataset: 'coco', 
      dataset_options: [],
      param: '', 
      param_options: [],

      labelCol: { span: 6 },
      wrapperCol: { span: 18 },
      train_param: {
        model: 'yolov5s.yaml',
        epochs: 16,
        batch: 4,
        other: '--img 640',
      },

      train_code: '',
      train_loading: false, 
      log_loading: true, 

      confirm_visible: false, 
    }
  },
  methods: {
    onChange(current) {
      console.log('onChange:', current);
      this.current = current;
      this.show_error = false; 
      this.error_message = '';
      if(current == 1) this.getWeights();
      if(current == 0) this.getDatasets();
    },
    nextStep() {
      this.current = this.current + 1; 
      this.train_code = '$ python train.py --batch ' + this.train_param.batch + ' --epochs ' + this.train_param.epochs + ' --data '+ this.dataset + ' --weights ' + this.weight + ' ' + this.train_param.other
    }, 

    getDatasets() {
      if (this.dataset_options.length) return;
      this.$axios({
        url: this.yolov_datasets_url,
        method: 'get',
      }).then(res =>{
        console.log(res.data)
        res.data.forEach(dataset => {
            this.dataset_options.push({ value: dataset.name });
          });
      }).catch(err =>{
        console.log(err);
      })
    },
    changeDataset() {
      console.log(this.dataset);
      // this.yolov_detect_url = "http://127.0.0.1:5000/yolov/detect?dataset=" + this.weight;
      // console.log(this.yolov_detect_url)
    }, 
    checkLables() {
      this.$axios.get("/yolov/check_labels?dataset=" + this.dataset)
      .then(res => {
        console.log(res.data);
        if (res.data.dataset_length != res.data.label_length) {
          this.show_error = true;
          this.error_message = '数据集 "' + this.dataset + '" 图片与标注不匹配，请标注完全！'; 
        }
        else this.current++; 
      }).catch(err => {
        console.log(err)
      })
    }, 

    getWeights() {
      if (this.weight_options.length) return;
      this.$axios.get(this.yolov_weights_url)
      .then(res => {
        console.log(res.data);
        res.data.forEach(weight => {
          this.weight_options.push({ value: weight });
        });
      }).catch(err => {
        console.log(err)
      })
    }, 
    changeWeight() {
      console.log(this.weight);
      // this.yolov_detect_url = "http://127.0.0.1:5000/yolov/detect?weight=" + this.weight;
      // console.log(this.yolov_detect_url)
    }, 

    checkTrain() {
      this.$axios({
        method: 'get', 
        url: this.$axios.defaults.baseURL + '/yolov/preTrain', 
        params: {
           dataset: this.dataset 
        }
      }).then(res => {
        console.log(res.data);
        if(!res.data.exists) 
        {
          this.confirm_visible = false; 
          this.startTrain();  // 不存在训练记录
        }
        else {
          this.confirm_visible = true;  
        }
        
      }).catch(err => {});
    }, 

    startTrain() {
      this.confirm_visible = false; 
      this.log_loading = false;
      this.train_loading = true;
      document.getElementById('log-content').innerHTML = '';
      
      this.$axios({
        url: this.$axios.defaults.baseURL + '/yolov/train',
        method: "post", 
        data: qs.stringify({
          dataset: this.dataset, 
          weight: this.weight, 
          model: this.train_param.model, 
          epochs: this.train_param.epochs, 
          batch: this.train_param.batch, 
          other: this.train_param.other, 
        }),
      }).then(res => {
        console.log(res.data);
        this.train_loading = false; 
        // setTimeout(() => {
          this.$router.push('/yolov_view'); 
        // }, 3); 
      }).catch(err => {
        console.log(err)
      }); 
      // this.$router.push('/yolov_log'); 
    }
  },
  mounted() {
    this.getDatasets(); 
    this.train_code = '$ python train.py --batch ' + this.train_param.batch + ' --epochs ' + this.train_param.epochs + ' --data '+ this.dataset + ' --weights ' + this.weight + ' ' + this.train_param.other
    let socket = io(this.$axios.defaults.baseURL + '/test'); 
    socket.on('response', response => {
      console.log(response); 
      document.getElementById('log-content').append(response.data)
      document.getElementById('log-content').append("\r\n");
    });
  }
}
</script>

<style scoped>
  .train {
    /* padding: 16px;  */
  }

  .steps {
    padding: 20px 96px; 
    margin-bottom: 48px;
  }

  .steps-content {
    width: 360px; 
    margin: auto;
    /* text-align: center;  */
  }

  .steps-text {
    margin: 20px 0; 
    text-align: left;
    font-size: 14.5px;
    color: #333;
    line-height: 20px;
  }

  .steps-input {
    padding: 10px 0; 
  }

  .step-next {
    margin: 20px auto; 
    width: 100%; 
  }

  .error {
    margin: 0; 
    color: #f00;
    text-align: left; 
    font-size: 14.5px;
  }

  .train-content {
    width: 960px;
  }

  /deep/ .ant-page-header-heading-sub-title, /deep/ .ant-steps-item-description {
     font-size: 13.5px; 
  }

</style>